#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from std_msgs.msg import Float64MultiArray, String
import argparse
import time
import numpy as np
import os
from ament_index_python.packages import get_package_share_directory


def load_data(path):
    data = np.loadtxt(path, delimiter=",")
    if data.shape[1] != 7:
        raise ValueError(f"Data shape is {data.shape}, expected N×7.")
    return data


class FollowTrajCommander(Node):
    def __init__(self, path):
        super().__init__('follow_traj_commander')

        self.publisher_move_joint = self.create_publisher(Float64MultiArray, '/command_move_joint', 10)
        self.publisher_follow_traj = self.create_publisher(String, '/command_move_traj', 10)

        package_path = get_package_share_directory('zorci')
        default_path = os.path.join(package_path, 'data', 'follow_trajectory_template.txt')

        if path is None:
            path = default_path

        try:
            if not os.path.exists(path):
                raise FileNotFoundError(f"File {path} not found")
            data = load_data(path)
            self.traj_path = path
            self.get_logger().info(f"Loaded data from {path}")
        except Exception as e:
            self.get_logger().warn(f"Failed to load {path}: {e}")
            data = load_data(default_path)
            self.traj_path = default_path
            self.get_logger().info(f"Loaded data from default file {default_path}")

        self.move_joint_cmd = data[0, :].tolist()

    def wait_for_subscribers(self, timeout=5.0):
        self.get_logger().info("Waiting for subscribers to connect...")
        start_time = time.time()
        while (self.publisher_move_joint.get_subscription_count() == 0 or
               self.publisher_follow_traj.get_subscription_count() == 0):
            if time.time() - start_time > timeout:
                self.get_logger().warn("Timeout waiting for subscribers!")
                return False
            rclpy.spin_once(self, timeout_sec=0.1)
            time.sleep(0.1)
        self.get_logger().info("Subscribers connected.")
        return True

    def send_commands(self):
        msg_joint = Float64MultiArray()
        msg_joint.data = self.move_joint_cmd
        self.publisher_move_joint.publish(msg_joint)
        self.get_logger().info("Published first point to /command_move_joint")

        # 等待5秒同时处理事件循环，确保消息发送
        start = time.time()
        while time.time() - start < 5.0:
            rclpy.spin_once(self, timeout_sec=0.1)

        msg_traj = String()
        msg_traj.data = self.traj_path
        self.publisher_follow_traj.publish(msg_traj)
        self.get_logger().info("Published traj file path to /command_move_traj")

        # 再等待1秒确保消息发送
        start = time.time()
        while time.time() - start < 1.0:
            rclpy.spin_once(self, timeout_sec=0.1)


def main():
    parser = argparse.ArgumentParser(description='Send follow trajectory command.')
    parser.add_argument('trajectory_path', type=str, nargs='?', default=None,
                        help='Trajectory file path to command (optional).')
    args = parser.parse_args()

    rclpy.init()
    node = FollowTrajCommander(args.trajectory_path)
    try:
        if node.wait_for_subscribers():
            node.send_commands()
        else:
            node.get_logger().warn("No subscribers connected, aborting publish.")
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
